class WGAN(models.Model):
    def __init__(self):
        super(WGAN, self).__init__()
        self.generator = models.Sequential([
            layers.Dense(128, activation='relu', input_dim=100),
            layers.Dense(784, activation='tanh')
        ])
        self.discriminator = models.Sequential([
            layers.Dense(128, activation='relu', input_dim=784),
            layers.Dense(1, activation='linear')
        ])

    def compile(self, g_optimizer, d_optimizer):
        super(WGAN, self).compile()
        self.g_optimizer = g_optimizer
        self.d_optimizer = d_optimizer

    def train_step(self, real_images):
        batch_size = tf.shape(real_images)[0]
        noise = tf.random.normal([batch_size, 100])

        with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
            generated_images = self.generator(noise, training=True)
            real_output = self.discriminator(real_images, training=True)
            fake_output = self.discriminator(generated_images, training=True)

            gen_loss = -tf.reduce_mean(fake_output)
            disc_loss = tf.reduce_mean(fake_output) - tf.reduce_mean(real_output)

        gradients_of_generator = gen_tape.gradient(gen_loss, self.generator.trainable_variables)
        gradients_of_discriminator = disc_tape.gradient(disc_loss, self.discriminator.trainable_variables)

        self.g_optimizer.apply_gradients(zip(gradients_of_generator, self.generator.trainable_variables))
        self.d_optimizer.apply_gradients(zip(gradients_of_discriminator, self.discriminator.trainable_variables))